
#include <assert.h>
#include <stdlib.h>
#include <string.h>
#include "utility.h"
#include "bit-stream.h"

static int bit_stream_allocate(BitStream* bit_stream, int length);

BitStream* bit_stream_new()
{
	BitStream* bit_stream;

	bit_stream = (BitStream*)malloc(sizeof(BitStream));
	if(NULL == bit_stream)
	{
		RES_NULL("Invoke malloc error.");
	}
	bit_stream->_data = NULL;
	bit_stream->_length = 0;
	return bit_stream;
}

BitStream* bit_stream_new_from_num(int length, unsigned int num)
{
	BitStream* bit_stream;
	unsigned char* data;
	unsigned int mask;
	unsigned char* ptr;

	if(length <= 0 || length > (sizeof(unsigned int) << 3))
	{
		RES_NULL("Argument is invalid.");
	}

	bit_stream = bit_stream_new();
	if(NULL == bit_stream)
	{
		RES_NULL("Invoke bit_stream_new error.");
	}
	if(bit_stream_allocate(bit_stream, length) < 0)
	{
		bit_stream_destroy(bit_stream);
		RES_NULL("Invoke bit_stream_allocate error.");
	}
	data = bit_stream->_data;
	mask = 1;
	for(ptr = data + length - 1; ptr >= data; --ptr)
	{
		*ptr = (0 == (num & mask)) ? 0 : 1;
		mask = mask << 1;
	}
	return bit_stream;
}

BitStream* bit_stream_new_from_bytes(int size, unsigned char* bytes)
{
	BitStream* bit_stream;
	int i, j;
	unsigned int mask;
	unsigned char* ptr;

	if(size <= 0 || NULL == bytes)
	{
		RES_NULL("Argument is invalid.");
	}

	bit_stream = bit_stream_new();
	if(NULL == bit_stream)
	{
		RES_NULL("Invoke bit_stream_new error.");
	}
	if(bit_stream_allocate(bit_stream, size << 3) < 0)
	{
		bit_stream_destroy(bit_stream);
		RES_NULL("Invoke bit_stream_allocate error.");
	}
	ptr = bit_stream->_data;
	for(i = 0; i < size; ++i)
	{
		mask = 0x80;
		for(j = 0; j < 8; ++j)
		{
			*ptr = (0 == (bytes[i] & mask)) ? 0 : 1;
			++ptr;
			mask = mask >> 1;
		}
	}
	return bit_stream;
}

void bit_stream_destroy(BitStream* bit_stream)
{
	if(bit_stream != NULL)
	{
		if(bit_stream->_data != NULL)
		{
			free(bit_stream->_data);
		}
		free(bit_stream);
	}
}

int bit_stream_append(BitStream* bit_stream, BitStream* combined)
{
	unsigned char* data;

	if(NULL == bit_stream || NULL == combined)
	{
		RES_INVAL("Argument is invalid.");
	}
	if(combined->_length < 0)
	{
		RES_INVAL("Argument is invalid.");
	}

	if(0 == combined->_length)
	{
		return 0;
	}
	data = (unsigned char*)malloc(bit_stream->_length + combined->_length);
	if(NULL == data)
	{
		RES_INVAL("Invoke malloc error.");
	}
	if(bit_stream->_length != 0)
	{
		memcpy(data, bit_stream->_data, bit_stream->_length);
		free(bit_stream->_data);
	}
	memcpy(data + bit_stream->_length, combined->_data, combined->_length);
	bit_stream->_data = data;
	bit_stream->_length += combined->_length;
	return 0;
}

int bit_stream_append_num(BitStream* bit_stream, int length, unsigned int num)
{
	BitStream* combined;
	int result;

	if(NULL == bit_stream || length <= 0)
	{
		RES_INVAL("Argument is invalid.");
	}

	combined = bit_stream_new_from_num(length, num);
	if(NULL == combined)
	{
		RES_INVAL("Invoke bit_stream_new_from_num error.");
	}
	result = bit_stream_append(bit_stream, combined);
	bit_stream_destroy(combined);
	if(result < 0)
	{
		RES_INVAL("Invoke bit_stream_append error.");
	}
	return 0;
}

int bit_stream_append_bytes(BitStream* bit_stream, int size, unsigned char* bytes)
{
	BitStream* combined;
	int result;

	if(NULL == bit_stream || size <= 0 || NULL == bytes)
	{
		RES_INVAL("Argument is invalid.");
	}

	combined = bit_stream_new_from_bytes(size, bytes);
	if(NULL == combined)
	{
		RES_INVAL("Invoke bit_stream_new_from_bytes error.");
	}
	result = bit_stream_append(bit_stream, combined);
	bit_stream_destroy(combined);
	if(result < 0)
	{
		RES_INVAL("Invoke bit_stream_append error.");
	}
	return 0;
}

unsigned char* bit_stream_convert_to_bytes(BitStream* bit_stream)
{
	int size, bytes, i, j, remainder;
	unsigned char* result, * ptr;
	unsigned char value;

	if(NULL == bit_stream)
	{
		RES_NULL("Argument is invalid.");
	}
	if(NULL == bit_stream->_data || bit_stream->_length <= 0)
	{
		RES_NULL("Argument is invalid.");
	}
	
	size = (bit_stream->_length + 7) >> 3;
	result = (unsigned char*)malloc(size);
	if(NULL == result)
	{
		RES_NULL("Invoke malloc error.");
	}
	ptr = bit_stream->_data;
	bytes = bit_stream->_length >> 3;
	for(i = 0; i < bytes; ++i)
	{
		value = 0;
		for(j = 0; j < 8; ++j)
		{
			value = value << 1;
			value |= *ptr;
			++ptr;
		}
		result[i] = value;
	}
	if(size != bytes)
	{
		assert(bytes + 1 == size);
		value = 0;
		remainder = bit_stream->_length & 7;
		assert(remainder < 8);
		for(j = 0; j < remainder; ++j)
		{
			value = value << 1;
			value |= *ptr;
			++ptr;
		}
		result[bytes] = value;
	}
	return result;
}


int bit_stream_allocate(BitStream* bit_stream, int length)
{
	unsigned char* data;

	if(NULL == bit_stream || length <= 0)
	{
		RES_INVAL("Argument is invalid.");
	}

	data = (unsigned char*)malloc(length);
	if(NULL == data)
	{
		RES_INVAL("Invoke malloc error.");
	}
	if(bit_stream->_data != NULL)
	{
		free(bit_stream->_data);
	}
	bit_stream->_data = data;
	bit_stream->_length = length;
	return 0;
}

